"""
CIFAR-100数据集处理 (使用torchvision实现)
"""

import random
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR100
from torch.utils.data import Dataset

from config import cfg
from PIL import Image
from config import cfg

class CIFAR100Dataset(Dataset):
    """CIFAR-100数据集 (使用torchvision实现)"""

    def __init__(self, root, train=True, known_classes=None, transform=None):
        self.root = root
        self.train = train
        self.transform = transform

        # 设置已知类（前50类）
        if known_classes is None:
            self.known_classes = list(range(cfg.NUM_KNOWN_CLASSES))
        else:
            self.known_classes = known_classes

        # 使用torchvision加载数据集
        self.dataset = CIFAR100(
            root=root,
            train=train,
            download=True,  # 自动下载数据集
        )

        # 获取原始数据和标签
        self.data = self.dataset.data
        self.labels = self.dataset.targets

        # 过滤数据
        if train:
            # 训练集只包含已知类
            indices = [
                i for i, label in enumerate(self.labels) if label in self.known_classes
            ]
            self.data = self.data[indices]
            self.labels = [self.labels[i] for i in indices]
            # 重新映射标签
            self.labels = [self.known_classes.index(label) for label in self.labels]
        else:
            # 测试集包含所有类，未知类标记为-1
            new_labels = []
            for label in self.labels:
                if label in self.known_classes:
                    new_labels.append(self.known_classes.index(label))
                else:
                    new_labels.append(-1)
            self.labels = new_labels

        # 统计
        self.num_known = sum(1 for label in self.labels if label >= 0)
        self.num_unknown = sum(1 for label in self.labels if label == -1)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        img, label = self.data[index], self.labels[index]

        # 转换为PIL图像
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        return img, label


class IncrementalCIFAR100Dataset(Dataset):
    """增量学习CIFAR-100数据集 (使用torchvision实现)"""

    def __init__(
        self, root, train=True, class_range=None, samples_per_class=0, transform=None
    ):
        self.root = root
        self.train = train
        self.transform = transform

        # 类别范围，例如[50, 60]表示类别50-59
        if class_range is None:
            self.class_range = list(range(100))
        else:
            self.class_range = list(range(class_range[0], class_range[1]))

        # 每类样本数限制（用于Few-shot）
        self.samples_per_class = samples_per_class

        # 使用torchvision加载数据集
        self.dataset = CIFAR100(
            root=root,
            train=train,
            download=True,  # 自动下载数据集
        )

        # 获取原始数据和标签
        self.data = self.dataset.data
        self.labels = self.dataset.targets

        # 过滤数据 - 只保留指定范围的类别
        indices = [
            i for i, label in enumerate(self.labels) if label in self.class_range
        ]
        self.data = self.data[indices]
        self.labels = [self.labels[i] for i in indices]

        # Few-shot采样
        if train and self.samples_per_class > 0:
            self._few_shot_sampling()

        # 统计每个类的样本数
        self.class_counts = {}
        for label in self.labels:
            self.class_counts[label] = self.class_counts.get(label, 0) + 1

        if not train:
            # 创建未知类标签：将不在class_range中的类标记为-1
            new_labels = []
            for label in self.labels:
                if label in self.class_range:
                    new_labels.append(label)
                else:
                    new_labels.append(-1)  # 未知类标记为-1
            self.labels = new_labels

    def _few_shot_sampling(self):
        """Few-shot采样：每个类只保留指定数量的样本"""
        # 按类别组织数据
        class_indices = {}
        for idx, label in enumerate(self.labels):
            if label not in class_indices:
                class_indices[label] = []
            class_indices[label].append(idx)

        # 每个类随机采样指定数量
        selected_indices = []
        for label, indices in class_indices.items():
            if len(indices) > self.samples_per_class:
                # 随机选择samples_per_class个样本
                selected = random.sample(indices, self.samples_per_class)
            else:
                # 如果样本不足，使用所有样本
                selected = indices
            selected_indices.extend(selected)

        # 更新数据
        selected_indices = sorted(selected_indices)
        self.data = self.data[selected_indices]
        self.labels = [self.labels[i] for i in selected_indices]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        img, label = self.data[index], self.labels[index]

        # 转换为PIL图像
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        return img, label


class ReplayBuffer:
    """回放缓冲区，用于存储旧类样本"""

    def __init__(self, buffer_size, num_classes):
        self.buffer_size = buffer_size
        self.num_classes = num_classes
        self.samples_per_class = buffer_size // num_classes

        self.data = []
        self.labels = []

    def update(self, dataset, model=None, device=None):
        """从数据集中选择样本更新缓冲区"""
        # 按类别组织数据
        class_indices = {}
        for idx in range(len(dataset)):
            _, label = dataset[idx]
            if label not in class_indices:
                class_indices[label] = []
            class_indices[label].append(idx)

        # 清空缓冲区
        self.data = []
        self.labels = []

        # 每个类选择样本
        for label in range(self.num_classes):
            if label in class_indices:
                indices = class_indices[label]
                # 随机选择或使用特定策略选择样本
                if len(indices) > self.samples_per_class:
                    selected = random.sample(indices, self.samples_per_class)
                else:
                    selected = indices

                for idx in selected:
                    img, lbl = dataset[idx]
                    self.data.append(img)
                    self.labels.append(lbl)

    def get_data(self):
        """获取缓冲区中的数据"""
        return list(zip(self.data, self.labels))


def get_transforms():
    """获取数据变换"""

    if cfg.USE_DATA_AUGMENTATION:
        # 使用数据增强
        train_transform = transforms.Compose(
            [
                transforms.Resize(
                    cfg.INPUT_SIZE,
                    interpolation=transforms.InterpolationMode.BICUBIC,
                ),
                transforms.RandomCrop(cfg.INPUT_SIZE, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean=cfg.PRETRAINED_MEAN, std=cfg.PRETRAINED_STD),
            ]
        )
    else:
        # 不使用数据增强（论文发现短期训练时效果更好）
        train_transform = transforms.Compose(
            [
                transforms.Resize(
                    cfg.INPUT_SIZE,
                    interpolation=transforms.InterpolationMode.BICUBIC,
                ),
                transforms.CenterCrop(cfg.INPUT_SIZE),
                transforms.ToTensor(),
                transforms.Normalize(mean=cfg.PRETRAINED_MEAN, std=cfg.PRETRAINED_STD),
            ]
        )

    # 测试数据预处理
    test_transform = transforms.Compose(
        [
            transforms.Resize(
                cfg.INPUT_SIZE, interpolation=transforms.InterpolationMode.BICUBIC
            ),
            transforms.CenterCrop(cfg.INPUT_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(mean=cfg.PRETRAINED_MEAN, std=cfg.PRETRAINED_STD),
        ]
    )

    return train_transform, test_transform
